import yaml
import pandas as pd
import os

def generate_config(servers, output_file='config.yaml'):
    """
    根据服务器表格数据生成 YAML 配置文件
    :param servers: 服务器配置信息列表
    :param output_file: 输出的配置文件名
    """
    config = {'servers': servers}
    
    # 将配置数据写入 YAML 文件
    with open(output_file, 'w') as file:
        yaml.dump(config, file, default_flow_style=False, sort_keys=False)
    print(f"Configuration file '{output_file}' generated successfully.")

def load_servers(file_path, sheet_name=0, delimiter=','):
    """
    根据文件类型 (CSV 或 XLSX) 加载服务器配置数据，并支持多地址处理
    :param file_path: 文件路径 (可以是 CSV 或 XLSX)
    :param sheet_name: 如果是 XLSX 文件，工作表名称或索引，默认第一个工作表
    :param delimiter: 用于分隔 src 和 dst 列中多个地址的分隔符，默认使用逗号
    :return: 服务器配置信息列表
    """
    file_ext = os.path.splitext(file_path)[1].lower()
    
    # 根据文件类型读取数据
    if file_ext == '.csv':
        df = pd.read_csv(file_path)
    elif file_ext in ['.xls', '.xlsx']:
        df = pd.read_excel(file_path, sheet_name=sheet_name)
    else:
        raise ValueError("Unsupported file format. Please provide a CSV or XLSX file.")
    
    servers = []

    for _, row in df.iterrows():
        # 分割多个地址，并去除可能存在的多余空格
        src_list = [src.strip() for src in row['src'].split(delimiter)]
        dst_list = [dst.strip() for dst in row['dst'].split(delimiter)]
        dst_ports = list(map(int, row['dstport'].split(delimiter)))  # 处理多个目的端口
        
        server = {
            'src': src_list,
            'dst': dst_list,
            'dstport': dst_ports,
            'description': row['description'],
            'ssh_user': row['ssh_user'],
            'ssh_key': row['ssh_key'],
        }
        # 如果 ssh_port 列存在且有值，添加 ssh_port，否则跳过
        if pd.notna(row.get('ssh_port', None)):
            server['ssh_port'] = int(row['ssh_port'])
        
        servers.append(server)
    
    return servers

if __name__ == '__main__':
    # 输入文件路径 (CSV 或 XLSX)
    file_path = 'servers.xlsx'  # 可以替换为 'servers.csv'

    # 从文件加载服务器数据
    servers_data = load_servers(file_path, delimiter=',')

    # 生成 YAML 配置文件
    generate_config(servers_data)
